#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "device_util.h"

using namespace optix;

rtDeclareVariable(rtObject,      top_object, , );
rtDeclareVariable(float3, shading_normal, attribute shading_normal, ); 
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, ); 
rtDeclareVariable(float,      t_hit,        rtIntersectionDistance, );
rtDeclareVariable(optix::Ray, ray,          rtCurrentRay, );

rtDeclareVariable(uint2, launch_index, rtLaunchIndex, );
rtBuffer<float4, 2>   result_buffer;
rtBuffer<ViewRayData, 2>  hit_buffer;

rtDeclareVariable(ViewRayData, ray_payload, rtPayload, );
rtDeclareVariable(uint2, launch_dim,   rtLaunchDim, );

rtDeclareVariable(float3,        eye, , );
rtDeclareVariable(float3,        U, , );
rtDeclareVariable(float3,        V, , );
rtDeclareVariable(float3,        W, , );

rtDeclareVariable(float3,  Ks, , );
rtDeclareVariable(float3,  Kd, , );
rtDeclareVariable(float3,  emissive_color, , );

RT_PROGRAM void closest_hit()
{ 
	float3 direction    = ray.direction;
	float3 origin       = ray.origin;
	float3 world_shading_normal   = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, shading_normal ) );
	float3 world_geometric_normal = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, geometric_normal ) );
	float3 ffnormal     = faceforward( world_shading_normal, -direction, world_geometric_normal );
	float3 hit_point    = origin + t_hit*direction;

	  // Check if this is a light source
  if( fmaxf( emissive_color ) > 0.0f ) {
    //ray_payload.position = hit_point;
		//ray_payload.normal = ffnormal;
		ray_payload.color = make_float3(1000.f, 1000.f, 1000.f);
    return;
  }
	/*
	if (fmaxf(Ks) > 0.f) {
		ray_payload.color = ray_payload.color * Ks; 
    
		float3 R = reflect( direction, ffnormal );
    optix::Ray refl_ray( hit_point, R, 0, 1e-2f );
    rtTrace( top_object, refl_ray, ray_payload ); 
	}
	else */{ 
		ray_payload.position = hit_point;
		ray_payload.normal = ffnormal;
		ray_payload.color = ray_payload.color * Kd;
	}

	/*
		float3 n = geometric_normal;
	n.x = fabsf(n.x);
	n.y = fabsf(n.y);
	n.z = fabsf(n.z);
  ray_payload.color = n; */
}

RT_PROGRAM void cast() 
{
	float2 screen = make_float2( result_buffer.size() ); 

  float2 d = make_float2(launch_index) / make_float2(launch_dim) * 2.f - 1.f;
  float3 ray_origin = eye;
  float3 ray_direction = normalize(d.x*U + d.y*V + W);
	 
  optix::Ray ray(ray_origin, ray_direction, 0, 1e-3f);


  ViewRayData prd; 
  prd.depth   = 0u; 
	prd.normal = make_float3(0.f);
	prd.position = make_float3(0.f);
	prd.color = make_float3(1.f);
	prd.importance = 1.f;

  rtTrace( top_object, ray, prd );

	hit_buffer[launch_index] = prd;
	result_buffer[launch_index] = make_float4(prd.color, 0.0f);
}

RT_PROGRAM void miss()
{
	ray_payload.color = make_float3(1.f, 1.f, 0.f);
}

RT_PROGRAM void except()
{
  const unsigned int code = rtGetExceptionCode();
  rtPrintf( "Caught exception 0x%X at launch index (%d,%d)\n", code, launch_index.x, launch_index.y );

	result_buffer[launch_index] = make_float4(1.f, 0.f, 1.f, 0.f);
}

RT_PROGRAM void except1()
{
	const unsigned int code = rtGetExceptionCode();
  rtPrintf( "Caught exception 0x%X at launch index (%d,%d)\n", code, launch_index.x, launch_index.y );
	result_buffer[launch_index] = make_float4(.5f, 0.f, .5f, 0.f);
}